package middleware.optimize;

import middleware.llvm.BasicBlock;
import middleware.llvm.instruction.AllocaInstruction;
import middleware.llvm.instruction.BrInstruction;
import middleware.llvm.instruction.LoadInstruction;
import middleware.llvm.instruction.StoreInstruction;
import middleware.llvm.value.FunctionArgumentIRType;
import middleware.llvm.value.FunctionIRType;

import java.util.*;
import java.util.stream.IntStream;

/**
 * 活跃变量分析，用于确定变量在程序中的活跃范围
 *
 * @author 杨凯伟
 * @date 2024/10/30
 */
public class LiveVariableAnalyzer {
    private final Map<BasicBlock, Set<AllocaInstruction>> defSets = new HashMap<>();
    private final Map<BasicBlock, Set<AllocaInstruction>> inSets = new HashMap<>();
    private final Map<BasicBlock, Set<AllocaInstruction>> outSets = new HashMap<>();
    private final Map<BasicBlock, List<BasicBlock>> dataFlowDiagram = new HashMap<>();
    private final List<BasicBlock> basicBlocks = new ArrayList<>();
    private boolean updateWithIterate = true;

    public LiveVariableAnalyzer(FunctionIRType func) {
        basicBlocks.addAll(func.getBasicBlocks());

        buildDataflowDiagram(func);
        func.getBasicBlocks().forEach(block -> {
            defSets.put(block, getDefSet(block));
            inSets.put(block, getUseSet(block)); // use set is useless, just use in set instead
            outSets.put(block, new HashSet<>());
        });
    }

    private void buildDataflowDiagram(FunctionIRType func) {
        func.getBasicBlocks().forEach(block -> dataFlowDiagram.put(block, getNextBlocks(block)));
    }

    private List<BasicBlock> getNextBlocks(BasicBlock block) {
        List<BasicBlock> nextBlocks = new ArrayList<>();
        var instructions = block.getInstructions();
        var lastInst = instructions.get(instructions.size() - 1);

        if (lastInst instanceof BrInstruction brInst) {
            if (brInst.getDest() != null) {
                nextBlocks.add(brInst.getDest());
            } else {
                nextBlocks.add(brInst.getTrueBranch());
                nextBlocks.add(brInst.getFalseBranch());
            }
        }
        return nextBlocks;
    }

    private Set<AllocaInstruction> getUseSet(BasicBlock block) {
        Set<AllocaInstruction> useSet = new HashSet<>();

//        for (var inst : block.getInstructions()) {
//            if (inst instanceof LoadInstruction loadInst) {
//                var ptr = loadInst.getPtr();
//                if (ptr instanceof AllocaInstruction allocaInstPtr) {
//                    useSet.add(allocaInstPtr);
//                }
//            }
//        }
        block.getInstructions().forEach(inst -> {
            if (inst instanceof LoadInstruction loadInst) {
                var ptr = loadInst.getPtr();
                if (ptr instanceof AllocaInstruction allocaInstPtr) {
                    useSet.add(allocaInstPtr);
                }
            }
        });

        return useSet;
    }

    private Set<AllocaInstruction> getDefSet(BasicBlock block) {
        Set<AllocaInstruction> defSet = new HashSet<>();

//        for (var inst : block.getInstructions()) {
//            if (inst instanceof StoreInstruction storeInst && !(storeInst.getValue() instanceof FunctionArgumentIRType)) {
//                var ptr = storeInst.getPtr();
//                if (ptr instanceof AllocaInstruction allocaInstPtr) {
//                    defSet.add(allocaInstPtr);
//                }
//            }
//        }
        block.getInstructions().forEach(inst -> {
            if (inst instanceof StoreInstruction storeInst && !(storeInst.getValue() instanceof FunctionArgumentIRType)) {
                var ptr = storeInst.getPtr();
                if (ptr instanceof AllocaInstruction allocaInstPtr) {
                    defSet.add(allocaInstPtr);
                }
            }
        });

        return defSet;
    }

    public void analyze() {
        while (!canStop()) {
            updateWithIterate = false;
            iterate();
        }
    }

    private boolean canStop() {
        return !updateWithIterate;
    }

    private void iterate() {
//        for (int i = basicBlocks.size() - 1; i >= 0; i--) {
//            var block = basicBlocks.get(i);
//            updateOutSet(block);
//            updateInSet(block);
//        }
        IntStream.range(0, basicBlocks.size()).map(i -> basicBlocks.size() - 1 - i).forEach(i -> {
            var block = basicBlocks.get(i);
            updateOutSet(block);
            updateInSet(block);
        });
    }

    private void updateOutSet(BasicBlock block) {
        var nextBlocks = dataFlowDiagram.get(block);

        var outSet = outSets.get(block);

//        for (var nextBlock : nextBlocks) {
//            int preSize = outSet.size();
//            outSet.addAll(inSets.get(nextBlock));
//            int nowSize = outSet.size();
//            if (preSize != nowSize) {
//                updateWithIterate = true;
//            }
//        }
        Iterator<BasicBlock> iterator = nextBlocks.iterator();
        boolean updated = false;
        while (iterator.hasNext()) {
            BasicBlock nextBlock = iterator.next();
            int preSize = outSet.size();
            outSet.addAll(inSets.get(nextBlock));
            int nowSize = outSet.size();
            if (preSize != nowSize) {
                updated = true;
            }
        }
        updateWithIterate = updated;
    }

    private void updateInSet(BasicBlock block) {
        var inSet = inSets.get(block);
        var outSet = outSets.get(block);
        var defSet = defSets.get(block);

        Set<AllocaInstruction> outSetCopy = new HashSet<>(outSet);
        outSetCopy.removeAll(defSet);

        inSet.addAll(outSetCopy);//outWithoutDef
    }

    public Map<BasicBlock, Set<AllocaInstruction>> getInSets() {
        return inSets;
    }

    public Map<BasicBlock, Set<AllocaInstruction>> getOutSets() {
        return outSets;
    }

    public Map<BasicBlock, Set<AllocaInstruction>> getDefSets() {
        return defSets;
    }
}
